import os
import re
from string import whitespace

import cv2
import paddlehub as hub


def imgdeleteBg(path, path_output):
    # 加载模型
    humanseg = hub.Module(name='deeplabv3p_xception65_humanseg')
    # humanseg = hub.Module(name='humanseg_server')
    # 获取文件列表
    if not os.path.exists(path_output):
        os.makedirs(path_output)
        print('make output_dir:' + path_output)
    for dirpath, dirnames, filenames in os.walk(path):
        for filename in filenames:
            try:
                if filename.__contains__(".jpg") or filename.__contains__(".png"):
                    imgpath = os.path.join(dirpath, filename)
                    # deeplabv3p_xception65_humanseg调用
                    results = humanseg.segmentation(images=[cv2.imread(imgpath)],
                                                    use_gpu=False,
                                                    visualization=True,
                                                    output_dir=path_output)
                    # # humanseg_server调用
                    # results = humanseg.segment(images=[cv2.imread(imgpath)],
                    #                                 use_gpu=False,
                    #                                 visualization=True,
                    #                                 output_dir=path_output)
                    for result in results:
                        print(result)
            except Exception as e:
                print(e)
                continue


def videodeleteBg(path, path_output):
    # 加载模型
    humanseg = hub.Module(name='humanseg_server')
    # 获取文件列表
    if not os.path.exists(path_output):
        os.makedirs(path_output)
        print('make output_dir:' + path_output)
    for dirpath, dirnames, filenames in os.walk(path):
        for filename in filenames:
            try:
                v_path = os.path.join(dirpath, filename)
                v_outpath = os.path.join(path_output, filename)
                results = humanseg.video_segment(video_path=v_path,
                                                 save_dir=v_outpath)
                for result in results:
                    print(result)
            except Exception as e:
                print(e)
                continue

def rename(file_path):
    for dirpath, dirnames, filenames in os.walk(file_path):
        for filename in filenames:
            try:
                inputfilename = dirpath.replace("\\", "/") + '/' + filename
                tables = str.maketrans('', '', "()【】-  */" + r"""!"#$%&'()*+，,-/:;<=>?@[\]^_`{|}~""" + whitespace)
                # 去除中文
                pattern = re.compile(u"[\u4e00-\u9fa5]+")
                filename = re.sub(pattern, "", filename)
                outfilename = dirpath.replace("\\", "/") + '/' + str.strip(filename.translate(tables))
                print('inputfilename:' + inputfilename)
                print('outfilename:' + outfilename)
                if not inputfilename.__eq__(outfilename) and os.path.exists(outfilename):
                    os.remove(outfilename)
                os.rename(inputfilename, outfilename)
            except FileExistsError as f:
                print(f)
                continue

if __name__ == "__main__":
    FrameCut_Path = r'D:\like\picture\202303'
    FrameSeg_Path = r'D:\like\picture_deleteBg\202303'
    type = 1  # 1 image  2 video
    rename(FrameCut_Path)
    if 1 == type:
        imgdeleteBg(FrameCut_Path, FrameSeg_Path)
    if 2 == type:
        videodeleteBg(FrameCut_Path, FrameSeg_Path)
